module mod_ode_solver_abstract

   use iso_fortran_env, only: dp => real64
   use petsc_wrapper

   private
   public :: type_ode_solver_abstract

   type :: type_rhs_func
      #: for i in range(1, 5)
      procedure(rhs_func_array${i}$d), pointer, nopass :: array${i}$d => null()
      #: endfor
      procedure(rhs_func_vector), pointer, nopass :: vector => null()
   end type type_rhs_func

   type, abstract :: type_ode_solver_abstract

      character(len=:), allocatable :: name
      character(len=:), allocatable :: scheme
      character(len=:), allocatable :: data_type
      integer :: order
      real(dp) :: dt, t, t_end
      #: for i, dimension in list(zip(range(1, 5), [":", ":, :", ":, :, :", ":, :, :, :"]))
      real(dp), pointer :: solution_array${i}$ (${dimension}$) => null()
      real(dp), allocatable :: solution_array${i}$_n(${dimension}$)
      real(dp), allocatable :: rhs_${i}$ (${dimension}$)
      #: endfor
      class(type_domain_vector), pointer :: solution_vector => null()
      type(type_domain_vector) :: solution_vector_n
      type(type_domain_vector) :: rhs_vector

      type(type_rhs_func) :: rhs_func

      type(type_domain_matrix), pointer :: A_vector => null()
      real(dp), pointer :: A_scalar(:, :) => null()

   contains

      procedure(initialize), deferred :: initialize
      procedure(destroy), deferred :: destroy

   generic :: step => &
                         #: for i in range(1, 5)
      step_array${i}$D, &
                         #:endfor
      step_vector
   
      generic :: substep => &
                         #: for i in range(1, 5)
      substep_array${i}$D, &
                         #:endfor
      substep_vector


   #: for i in range(1, 5)
   procedure(step_array${i}$D), deferred :: step_array${i}$D
   procedure(substep_array${i}$D), deferred :: substep_array${i}$D
   #: endfor
   procedure(step_vector), deferred :: step_vector
   procedure(substep_vector), deferred :: substep_vector

   generic :: set_solution => &
      #: for i in range(1, 5)
      SetSolution_array${i}$D, &
      #: endfor
      SetSolution_vector
   #: for i in range(1, 5)
   procedure, private :: SetSolution_array${i}$D
   #: endfor
   procedure, private :: SetSolution_vector

   procedure :: set_time_step => SetTimeStep
   procedure :: get_time_step => GetTimeStep

   procedure :: set_rhs_function => SetRHSFunction
   procedure :: set_rhs_jacobian => SetRHSJacobian

   generic :: solve_rhs => &
                              #: for i in range(1, 5)
      SolveRHS_array${i}$D, &
                              #: endfor
      SolveRHS_vector

   #: for i in range(1, 5)
   procedure, private :: SolveRHS_array${i}$D
   #: endfor
   procedure, private :: SolveRHS_vector

   end type type_ode_solver_abstract

   abstract interface
      subroutine initialize(this, scheme)
         import
         class(type_ode_solver_abstract), intent(inout) :: this
         character(len=*), intent(in), optional :: scheme
      end subroutine initialize

      subroutine destroy(this)
         import
         class(type_ode_solver_abstract), intent(inout) :: this
      end subroutine destroy

      #: for i, dimension in list(zip(range(1, 5), [":", ":, :", ":, :, :", ":, :, :, :"]))
      Subroutine step_array${i}$D(this, solution, ctx)
         import
         class(type_ode_solver_abstract), intent(inout) :: this
         real(dp), intent(inout), target :: solution(${dimension}$)
         type(c_ptr), intent(inout) :: ctx
      end subroutine step_array${i}$D
      Subroutine substep_array${i}$D(this, lrk, solution, ctx)
         import
         class(type_ode_solver_abstract), intent(inout) :: this
         real(dp), intent(inout), target :: solution(${dimension}$)
         type(c_ptr), intent(inout) :: ctx
         integer, intent(in) :: lrk
      end subroutine substep_array${i}$D
      #: endfor

      subroutine step_vector(this, solution, ctx)
         import
         class(type_ode_solver_abstract), intent(inout) :: this
         class(type_domain_vector), intent(inout), target :: solution
         type(c_ptr), intent(inout) :: ctx
      end subroutine step_vector

      subroutine substep_vector(this, lrk, solution, ctx)
         import
         class(type_ode_solver_abstract), intent(inout) :: this
         class(type_domain_vector), intent(inout), target :: solution
         type(c_ptr), intent(inout) :: ctx
         integer, intent(in) :: lrk
      end subroutine substep_vector

      #: for i, dimension in list(zip(range(1, 5), [":", ":, :", ":, :, :", ":, :, :, :"]))
      subroutine rhs_func_array${i}$d(u, rhs, t, ctx)
         import
         real(dp), intent(in), target :: u(${dimension}$)
         real(dp), intent(inout), target :: rhs(${dimension}$)
         real(dp), intent(in) :: t
         type(c_ptr), intent(inout) :: ctx

      end subroutine rhs_func_array${i}$d
      #: endfor
      subroutine rhs_func_vector(u, rhs, t, ctx)
         import
         class(type_domain_vector), intent(in), target :: u
         class(type_domain_vector), intent(inout), target :: rhs
         real(dp), intent(in) :: t
         type(c_ptr), intent(inout) :: ctx

      end subroutine rhs_func_vector
   end interface

   contains

   subroutine SetTimeStep(this, dt)
      class(type_ode_solver_abstract), intent(inout) :: this
      real(dp), intent(in) :: dt
      this%dt = dt
   end subroutine SetTimeStep

   function GetTimeStep(this) result(dt)
      class(type_ode_solver_abstract), intent(in) :: this
      real(dp) :: dt
      dt = this%dt
   end function GetTimeStep

   subroutine SetRHSFunction(this, array1d_function, array2d_function, array3d_function, vector_function)
      class(type_ode_solver_abstract), intent(inout) :: this
      procedure(rhs_func_array1d), optional :: array1d_function
      procedure(rhs_func_array2d), optional :: array2d_function
      procedure(rhs_func_array3d), optional :: array3d_function
      procedure(rhs_func_vector), optional :: vector_function
      if (present(array1d_function)) this%rhs_func%array1D => array1d_function
      if (present(array2d_function)) this%rhs_func%array2D => array2d_function
      if (present(array3d_function)) this%rhs_func%array3D => array3d_function
      if (present(vector_function)) this%rhs_func%vector => vector_function
   end subroutine SetRHSFunction

   subroutine SetRHSJacobian(this, A_scalar, A_vector)
      class(type_ode_solver_abstract), intent(inout) :: this
      real(dp), optional, dimension(:, :), target :: A_scalar
      type(type_domain_matrix), optional, target :: A_vector

      if(present(A_scalar)) this%A_scalar => A_scalar
      if(present(A_vector)) this%A_vector => A_vector

   end subroutine SetRHSJacobian

   #: for i, dimension in list(zip(range(1, 5), [":", ":, :", ":, :, :", ":, :, :, :"]))
   Subroutine SetSolution_array${i}$D (this, solution)
      class(type_ode_solver_abstract), intent(inout) :: this
      real(dp), intent(in), target :: solution(${dimension}$)

      this%solution_array${i}$ => solution
      this%data_type = "array${i}$D"
      if(.not. allocated(this%solution_array${i}$_n)) allocate (this%solution_array${i}$_n, source=solution)
      if(.not. allocated(this%rhs_${i}$)) allocate (this%rhs_${i}$, source=solution)

   end subroutine SetSolution_array${i}$D

   Subroutine GetSolution_array${i}$D (this, solution)
      class(type_ode_solver_abstract), intent(in) :: this
      real(dp), intent(inout) :: solution(${dimension}$)

      solution = this%solution_array${i}$

   end subroutine GetSolution_array${i}$D

   subroutine SolveRHS_array${i}$D(this, t, flux, ctx)
      class(type_ode_solver_abstract), intent(inout) :: this
      real(dp), intent(inout) :: flux(${dimension}$)
      real(dp), intent(in) :: t
      type(c_ptr), intent(inout), optional :: ctx

      call this%rhs_func%array${i}$D(this%solution_array${i}$, this%rhs_${i}$, t, ctx)

   end subroutine SolveRHS_array${i}$D
   #: endfor

   Subroutine SetSolution_vector(this, solution)
      class(type_ode_solver_abstract), intent(inout) :: this
      class(type_domain_vector), intent(in), target :: solution

      this%solution_vector => solution
      this%data_type = "vector"

   end subroutine SetSolution_vector

   Subroutine GetSolution_vector(this, solution)
      class(type_ode_solver_abstract), intent(in) :: this
      class(type_domain_vector), intent(inout) :: solution

      call solution%copy(this%solution_vector)

   end subroutine GetSolution_vector

   subroutine SolveRHS_vector(this, t, flux, ctx)
      class(type_ode_solver_abstract), intent(inout) :: this
      class(type_domain_vector), intent(inout) :: flux
      real(dp), intent(in) :: t
      type(c_ptr), intent(inout) :: ctx

      if (.not. this%solution_vector%allocated) call this%solution_vector%duplicate(flux)
      if (.not. this%rhs_vector%allocated) call this%rhs_vector%duplicate(flux)
      call this%rhs_func%vector(this%solution_vector, this%rhs_vector, t, ctx)

   end subroutine SolveRHS_vector

end module mod_ode_solver_abstract
